package com.zk.storm;
//用于IP获取地理信息
import com.maxmind.geoip2.DatabaseReader;
import com.maxmind.geoip2.model.CityResponse;
import com.maxmind.geoip2.record.City;
import com.maxmind.geoip2.record.Country;
import com.maxmind.geoip2.record.Location;
import org.apache.storm.task.OutputCollector;
import org.apache.storm.task.TopologyContext;
import org.apache.storm.topology.OutputFieldsDeclarer;
import org.apache.storm.topology.base.BaseRichBolt;
import org.apache.storm.tuple.Fields;
import org.apache.storm.tuple.Tuple;
import org.apache.storm.tuple.Values;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;

import java.io.File;
import java.net.InetAddress;
import java.util.Map;

public class IpLocationBolt extends BaseRichBolt {
    private OutputCollector collector;
    private DatabaseReader geoIPReader;  // GeoIP2数据库读取器
    private ObjectMapper jsonMapper;     // JSON解析器

    // 数据库路径
    private static final String GEOIP_DB_PATH = "src/main/resources/GeoLite2-City.mmdb";

    @Override
    public void prepare(Map<String, Object> topoConf, TopologyContext context, OutputCollector collector) {
        this.collector = collector;
        this.jsonMapper = new ObjectMapper();

        try {
            // 初始化GeoIP2数据库读取器
            File databaseFile = new File(GEOIP_DB_PATH);
            this.geoIPReader = new DatabaseReader.Builder(databaseFile).build();
        } catch (Exception e) {
            throw new RuntimeException("GeoIP数据库初始化失败", e);
        }
    }

    @Override
    public void execute(Tuple input) {
        try {
            // 1. 从输入元组获取原始消息（JSON字符串）
            String originalMessage = input.getStringByField("value");
            // 2. 解析JSON获取IP字段
            JsonNode jsonNode = jsonMapper.readTree(originalMessage);
            String ipStr = jsonNode.get("ip").asText();
            // 3. 使用GeoIP2查询地理位置
            InetAddress ip = InetAddress.getByName(ipStr); //转化为IP格式内容
            CityResponse response = geoIPReader.city(ip);
            // 4. 提取关键地理位置信息（处理可能的空值）
            Country country = response.getCountry();
            String countryName = (country != null) ? country.getName() : "未知";
            // 5. 输出包含地理位置的新元组（原始消息+地理位置信息）
            collector.emit(new Values(
                    originalMessage,  // 原始消息用于后续处理
                    countryName
            ));
            // 确认消息处理成功
            collector.ack(input);
        } catch (Exception e) {
            // 处理异常（如IP格式错误、数据库查询失败）
            System.err.println("IP地理位置转换失败: " + e.getMessage());
            collector.fail(input);  // 通知Storm重新发送该消息
        }
    }

    @Override
    public void declareOutputFields(OutputFieldsDeclarer declarer) {
        // 声明输出字段：原始消息、国家
        declarer.declare(new Fields(
                "original_message",
                "country"
        ));
    }

    @Override
    public void cleanup() {
        // 关闭资源（Storm关闭时调用）
        if (geoIPReader != null) {
            try {
                geoIPReader.close();
            } catch (Exception e) {
                System.err.println("GeoIP数据库关闭异常: " + e.getMessage());
            }
        }
    }
}
